"""
图片生成服务
"""

import os
import time
import uuid
from typing import Optional, List
import httpx
from sqlalchemy.orm import Session
from app.models.image_generation import ImageGeneration
from app.models.conversation import Conversation, Message
from app.schemas.image_generation import ImageGenerationRequest, ImageAdjustmentRequest
from app.core.logging_config import get_logger
from app.core.config import settings

logger = get_logger(__name__)


class ImageGenerationService:
    """图片生成服务类"""
    
    def __init__(self):
        self.upload_dir = os.path.join("uploads", "images")
        os.makedirs(self.upload_dir, exist_ok=True)
    
    async def generate_image(
        self,
        db: Session,
        user_id: int,
        request: ImageGenerationRequest
    ) -> ImageGeneration:
        """
        生成图片
        
        Args:
            db: 数据库会话
            user_id: 用户ID
            request: 图片生成请求
            
        Returns:
            ImageGeneration: 生成的图片记录
        """
        try:
            # 验证对话是否存在且属于该用户
            conversation = db.query(Conversation).filter(
                Conversation.id == request.conversation_id,
                Conversation.user_id == user_id
            ).first()
            
            if not conversation:
                raise ValueError("对话不存在或无权限访问")
            
            # 如果是调整操作,获取原始图片信息
            original_prompt = request.prompt
            if request.original_image_id:
                original_image = db.query(ImageGeneration).filter(
                    ImageGeneration.id == request.original_image_id,
                    ImageGeneration.user_id == user_id
                ).first()
                
                if original_image:
                    # 将原始提示词与调整描述结合
                    original_prompt = f"{original_image.prompt}. 调整: {request.prompt}"
            
            # 调用图片生成API (这里以OpenAI DALL-E为例)
            image_url, image_path, metadata = await self._call_image_api(
                prompt=original_prompt,
                model_provider=request.model_provider or "openai",
                model_name=request.model_name or "dall-e-3",
                size=request.size or "1024x1024",
                quality=request.quality or "standard",
                style=request.style,
                n=request.n or 1
            )
            
            # 保存图片生成记录
            image_gen = ImageGeneration(
                conversation_id=request.conversation_id,
                user_id=user_id,
                prompt=original_prompt,
                original_image_id=request.original_image_id,
                image_url=image_url,
                image_path=image_path,
                model_provider=request.model_provider,
                model_name=request.model_name,
                size=request.size,
                quality=request.quality,
                style=request.style,
                image_metadata=metadata
            )
            
            db.add(image_gen)
            
            # 在对话中添加图片生成消息
            user_message = Message(
                conversation_id=request.conversation_id,
                role="user",
                content=f"生成图片: {request.prompt}",
                message_metadata={"type": "image_request", "prompt": request.prompt}
            )
            db.add(user_message)
            
            assistant_message = Message(
                conversation_id=request.conversation_id,
                role="assistant",
                content=f"我已经为您生成了图片",
                message_metadata={
                    "type": "image_response",
                    "image_id": None,  # 将在commit后更新
                    "image_url": image_url
                }
            )
            db.add(assistant_message)
            
            # 更新对话消息计数
            conversation.message_count += 2
            
            db.commit()
            db.refresh(image_gen)
            
            # 更新assistant消息的image_id
            assistant_message.message_metadata["image_id"] = image_gen.id
            db.commit()
            
            logger.info(f"图片生成成功: user_id={user_id}, image_id={image_gen.id}")
            return image_gen
            
        except Exception as e:
            db.rollback()
            logger.error(f"图片生成失败: {str(e)}", exc_info=True)
            raise
    
    async def adjust_image(
        self,
        db: Session,
        user_id: int,
        request: ImageAdjustmentRequest
    ) -> ImageGeneration:
        """
        调整已生成的图片
        
        Args:
            db: 数据库会话
            user_id: 用户ID
            request: 图片调整请求
            
        Returns:
            ImageGeneration: 调整后的图片记录
        """
        try:
            # 构建完整的图片生成请求
            gen_request = ImageGenerationRequest(
                conversation_id=request.conversation_id,
                original_image_id=request.original_image_id,
                prompt=request.adjustment_prompt,
                size=request.size or "1024x1024",
                quality=request.quality or "standard",
                style=request.style
            )
            
            return await self.generate_image(db, user_id, gen_request)
            
        except Exception as e:
            logger.error(f"图片调整失败: {str(e)}", exc_info=True)
            raise
    
    async def _call_image_api(
        self,
        prompt: str,
        model_provider: str,
        model_name: str,
        size: str,
        quality: str,
        style: Optional[str],
        n: int
    ) -> tuple[str, str, dict]:
        """
        调用图片生成API
        
        Args:
            prompt: 提示词
            model_provider: 模型提供商
            model_name: 模型名称
            size: 图片尺寸
            quality: 图片质量
            style: 图片风格
            n: 生成数量
            
        Returns:
            tuple: (image_url, image_path, metadata)
        """
        try:
            if model_provider == "openai":
                return await self._call_openai_dalle(prompt, model_name, size, quality, style, n)
            else:
                raise ValueError(f"不支持的模型提供商: {model_provider}")
                
        except Exception as e:
            logger.error(f"调用图片生成API失败: {str(e)}", exc_info=True)
            # 返回模拟数据用于测试
            return await self._mock_image_generation(prompt, size)
    
    async def _call_openai_dalle(
        self,
        prompt: str,
        model_name: str,
        size: str,
        quality: str,
        style: Optional[str],
        n: int
    ) -> tuple[str, str, dict]:
        """
        调用OpenAI DALL-E API生成图片
        """
        try:
            # 从环境变量或配置文件获取API密钥
            api_key = getattr(settings, 'OPENAI_API_KEY', None)
            if not api_key:
                logger.warning("未配置OpenAI API密钥,使用模拟数据")
                return await self._mock_image_generation(prompt, size)
            
            async with httpx.AsyncClient(timeout=60.0) as client:
                headers = {
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json"
                }
                
                payload = {
                    "model": model_name,
                    "prompt": prompt,
                    "n": n,
                    "size": size
                }
                
                if quality:
                    payload["quality"] = quality
                if style:
                    payload["style"] = style
                
                response = await client.post(
                    "https://api.openai.com/v1/images/generations",
                    headers=headers,
                    json=payload
                )
                
                response.raise_for_status()
                data = response.json()
                
                # 获取生成的图片URL
                image_url = data["data"][0]["url"]
                
                # 下载图片到本地
                image_path = await self._download_image(image_url)
                
                metadata = {
                    "revised_prompt": data["data"][0].get("revised_prompt"),
                    "response": data
                }
                
                return image_url, image_path, metadata
                
        except Exception as e:
            logger.error(f"调用OpenAI DALL-E API失败: {str(e)}", exc_info=True)
            return await self._mock_image_generation(prompt, size)
    
    async def _download_image(self, image_url: str) -> str:
        """
        从URL下载图片到本地
        
        Args:
            image_url: 图片URL
            
        Returns:
            str: 本地文件路径
        """
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.get(image_url)
                response.raise_for_status()
                
                # 生成唯一文件名
                filename = f"{uuid.uuid4()}.png"
                filepath = os.path.join(self.upload_dir, filename)
                
                # 保存图片
                with open(filepath, "wb") as f:
                    f.write(response.content)
                
                logger.info(f"图片已下载到本地: {filepath}")
                return filepath
                
        except Exception as e:
            logger.error(f"下载图片失败: {str(e)}", exc_info=True)
            return ""
    
    async def _mock_image_generation(self, prompt: str, size: str) -> tuple[str, str, dict]:
        """
        模拟图片生成(用于开发和测试)
        
        Args:
            prompt: 提示词
            size: 图片尺寸
            
        Returns:
            tuple: (image_url, image_path, metadata)
        """
        # 使用占位图服务作为模拟
        width, height = size.split('x')
        image_url = f"https://via.placeholder.com/{width}x{height}.png?text={prompt[:20]}"
        
        metadata = {
            "mock": True,
            "prompt": prompt,
            "size": size,
            "note": "这是模拟数据,请配置真实的图片生成API"
        }
        
        logger.info(f"使用模拟图片生成: {image_url}")
        return image_url, "", metadata
    
    def get_conversation_images(
        self,
        db: Session,
        user_id: int,
        conversation_id: int,
        skip: int = 0,
        limit: int = 50
    ) -> List[ImageGeneration]:
        """
        获取对话中的所有图片
        
        Args:
            db: 数据库会话
            user_id: 用户ID
            conversation_id: 对话ID
            skip: 跳过记录数
            limit: 返回记录数
            
        Returns:
            List[ImageGeneration]: 图片列表
        """
        return db.query(ImageGeneration).filter(
            ImageGeneration.user_id == user_id,
            ImageGeneration.conversation_id == conversation_id
        ).order_by(ImageGeneration.created_at.desc()).offset(skip).limit(limit).all()
    
    def get_image_by_id(
        self,
        db: Session,
        user_id: int,
        image_id: int
    ) -> Optional[ImageGeneration]:
        """
        根据ID获取图片
        
        Args:
            db: 数据库会话
            user_id: 用户ID
            image_id: 图片ID
            
        Returns:
            Optional[ImageGeneration]: 图片记录
        """
        return db.query(ImageGeneration).filter(
            ImageGeneration.id == image_id,
            ImageGeneration.user_id == user_id
        ).first()

